#include <Eigen/Eigen>
#include <bits/stdc++.h>
#include <geometry_msgs/PoseStamped.h>
#include <mavros_msgs/State.h>
#include <ros/ros.h>
#include <trajectory_msgs/JointTrajectoryPoint.h>

mavros_msgs::State current_state;
Eigen::Vector3d current_p;
ros::Publisher pva_pub;
double takeoff_h;

void stateCallback(const mavros_msgs::State::ConstPtr &msg) {
  current_state = *msg;
}

void positionCallback(const geometry_msgs::PoseStamped::ConstPtr &msg) {
  current_p << msg->pose.position.x, msg->pose.position.y, msg->pose.position.z;
}

void setPVA(Eigen::Vector3d p, Eigen::Vector3d v, Eigen::Vector3d a,
            double yaw = 0.0) {
  trajectory_msgs::JointTrajectoryPoint pva_setpoint;

  pva_setpoint.positions.push_back(p(0)); // x
  pva_setpoint.positions.push_back(p(1)); // y
  pva_setpoint.positions.push_back(p(2)); // z
  pva_setpoint.positions.push_back(yaw);

  pva_setpoint.velocities.push_back(v(0));
  pva_setpoint.velocities.push_back(v(1));
  pva_setpoint.velocities.push_back(v(2));

  pva_setpoint.accelerations.push_back(a(0));
  pva_setpoint.accelerations.push_back(a(1));
  pva_setpoint.accelerations.push_back(a(2));

  pva_pub.publish(pva_setpoint);
}

std::vector<Eigen::Vector3d> traj_plan(const Eigen::Vector3d &start_p,
                                       const Eigen::Vector3d &end_p,
                                       double rate, double v = 1.0) {
  double delt_t = 1.0 / rate;
  std::vector<Eigen::Vector3d> traj;
  double total_t = (end_p - start_p).norm() / v;
  for (int i = 0; i * delt_t <= total_t; i++) {
    Eigen::Vector3d pos = start_p + i * delt_t / total_t * (end_p - start_p);
    traj.push_back(pos);
  }
  return traj;
}

Eigen::Vector3d send_traj(const std::vector<Eigen::Vector3d> &traj, int &idx) {
  idx = std::min((int)(traj.size() - 1), idx);
  Eigen::Vector3d pos = traj[idx];
  idx = std::min((int)(traj.size() - 1), idx + 1);
  return pos;
}

int main(int argc, char **argv) {

  ros::init(argc, argv, "takeoff_landing");
  ros::NodeHandle nh("~");
  ros::Subscriber state_sub =
      nh.subscribe<mavros_msgs::State>("/mavros/state", 1, stateCallback);
  ros::Subscriber pose_sub = nh.subscribe<geometry_msgs::PoseStamped>(
      "/mavros/local_position/pose", 1, positionCallback);
  pva_pub =
      nh.advertise<trajectory_msgs::JointTrajectoryPoint>("/pva_setpoint", 1);

  const int LOOPRATE = 30;
  ros::Rate loop_rate(LOOPRATE);
  Eigen::Vector3d takeoff_p, landoff_p;
  double takeoff_h = 3.0;
  std::vector<Eigen::Vector3d> traj;
  int idx = 0;
  double start_hover_time = 0.0;
  bool start_takeoff = false;
  bool start_hover = false;
  bool start_landing = false;

  while (true) {
    Eigen::Vector3d send_p;
    if (current_state.mode != "OFFBOARD") {
      start_takeoff = false;
      start_hover = false;
      start_landing = false;
      traj.clear();
      idx = 0;
      takeoff_p = current_p;
      takeoff_p(2) = current_p(2) + takeoff_h;
      send_p = current_p;
    } else {
      if (!start_takeoff) {
        traj = traj_plan(current_p, takeoff_p, LOOPRATE);
        idx = 0;
        start_takeoff = true;
        ROS_INFO("start takeoff...");
      }
      if (start_takeoff && !start_hover &&
          (current_p - takeoff_p).norm() < 0.3) {
        start_hover_time = ros::Time::now().toSec();
        start_hover = true;
        ROS_INFO("start hovering...");
      }
      if (start_hover && !start_landing &&
          ros::Time::now().toSec() - start_hover_time > 5.0) {
        start_landing = true;
        landoff_p = current_p;
        landoff_p(2) = -10.0;
        traj = traj_plan(current_p, landoff_p, LOOPRATE);
        idx = 0;
        ROS_INFO("start landing...");
      }
      send_p = send_traj(traj, idx);
    }
    ROS_INFO_THROTTLE(1, "send_p:%f %f %f", send_p(0), send_p(1), send_p(2));
    setPVA(send_p, Eigen::Vector3d::Zero(), Eigen::Vector3d::Zero(), 0.0);
    loop_rate.sleep();
    ros::spinOnce();
  }

  return 0;
}